#include <rtl/base.h>
#include "wdt.h"
//#include "usb.h"
#include "ag.h"
#include "drivers/timer.h"
#include "drivers/usbdcore.h"
#include "drivers/simple/mmc.h"
#include "bootinfo.h"

#include <rtl/memory.h>
#include <rtl/string.h>
#include <rtl/print.h>
#include <rtl/heap.h>

#define GLOBAL_HEAP_POOL_SIZE 16384
UINT8 AgpGlobalHeapPool[GLOBAL_HEAP_POOL_SIZE];
PVOID AgpGlobalHeap = NULL;

#define LK_LOAD_ADDRESS (0x46000000 - 512)
#define LK_ENTRY_POINT (0x46000000)
#define TEE_LOAD_ADDRESS (0x446005C0)
#define MTK_DA_LOAD_ADDRESS (0x52000000)
#define PL_LOAD_ADDRESS (0x201000)
//#define ATF_LOAD_ADDRESS (0x1005C0)
#define ATF_LOAD_ADDRESS (0x44601000)
#define ATF_MAX_SIZE (PL_LOAD_ADDRESS - ATF_LOAD_ADDRESS)

#define BUFFER_SIZE 8192
static char Buffer[BUFFER_SIZE] ALIGNED(16);

VOID AgEnterBootROM();
VOID AgpRetrieveRamConsole();
VOID ResetBoard()
{
    WRITE32(MTK_WDT_RESTART, MTK_WDT_RESTART_KEY);
    UINT32 v = READ32(MTK_WDT_MODE);
#if 1
    v &= ~MTK_WDT_MODE_AUTO_RESTART;
#endif
    v &= ~(MTK_WDT_MODE_IRQ | MTK_WDT_MODE_DUAL_MODE);
#if 1
    v |= (MTK_WDT_MODE_KEY | MTK_WDT_MODE_EXTEN | MTK_WDT_MODE_AUTO_RESTART);
#else
    v |= (MTK_WDT_MODE_KEY | MTK_WDT_MODE_EXTEN);
#endif
    WRITE32(MTK_WDT_MODE, v);
    for (int i = 0; i < 500; i++)
    {
    }
    WRITE32(MTK_WDT_SWRST, MTK_WDT_SWRST_KEY);

    for (;;)
        ;
}

static BOOLEAN usb_connect(UINT32 tmo)
{
    ULONG start_time = get_timer(0);
    BOOLEAN result = FALSE;
    UINT32 i = 1;

    mt_usb_disconnect_internal();
    mt_usb_connect_internal();

#if CFG_USBIF_COMPLIANCE
    /* USB compliance test: 100mA charging current when USB is unconfigured. */
    platform_set_chrg_cur(70);
#endif

    do
    {
        /* kick wdt to avoid cpu reset during usb driver installation if not present */
        //platform_wdt_all_kick();
        service_interrupts();

        if (usbdl_configured())
        {
#if CFG_USBIF_COMPLIANCE
            /* USB compliance test: 500mA charging current when USB is configured but
             * we set the charging current to 450mA since 500mA doesn't support in the
             * platform.
             */
            platform_set_chrg_cur(450);
#endif
            result = TRUE;
            break;
        }

        if (tmo)
        {
            /* enable timeout mechanism */
            if (get_timer(start_time) > tmo)
                break;
#if !CFG_FPGA_PLATFORM
            /* cable plugged-out and power key detection each 1 second */
            if (get_timer(start_time) > i * 1000)
            {
                /* Stop booting when usb boot (ex. KPOC) but plug-out usb. */
                //if (!usb_accessory_in() && !pmic_detect_powerkey() && (g_boot_reason == BR_USB))
                //    pl_power_off();

#if !CFG_EVB_PLATFORM
                /* check bypass power key from the 2nd second */
                /*if (i > 1 && pmic_detect_powerkey())
                {
                    print("%s Enumeration(Skip): powerkey pressed\n", MOD);
                    break;
                }*/
                i++;
#endif
            }
#endif
        }
    } while (1);

    return result;
}

static void usb_disconnect(void)
{
    mt_usb_disconnect_internal();
}

extern char __bss_start[];
extern char __bss_end[];

BOOLEAN AgpConsoleEnabled = FALSE;
BOOLEAN AgpHaveLk = FALSE;
BOOLEAN AgpHaveMtkDa = FALSE;
BOOLEAN AgpHavePl = FALSE;
SIZE AgpDaSize = 0;
SIZE AgpLkSize = 0;
BOOLEAN AgpHaveAtf = FALSE;
SIZE AgpAtfSize = 0;

SIZE AgPrint(VOID *Data, CONST VOID *Buffer, SIZE Length)
{
    if (Length == 0 || !AgpConsoleEnabled)
        return Length;

    PACKET_HEADER Header;
    Header.Magic = PACKET_HEADER_MAGIC;
    Header.Type = PT_CONSOLE;
    Header.DataLength = Length;
    Header.Serial = 0;

    mt_usbtty_putcn(PACKET_HEADER_SIZE, (char *)&Header, 1);
    mt_usbtty_puts((char *)Buffer, Length);

    return Length;
}

VOID AgpSetup()
{
    AgpConsoleEnabled = FALSE;
    RtlSetOutputStreamHandler(AgPrint);

    UsbSendStatus(AG_STATUS_OK, NULL, 0);
}

VOID AgpReceivePacketHeader(PPACKET_HEADER Header)
{
    while (TRUE)
    {
        UINT8 Magic[2];
        Magic[0] = mt_usbtty_getc();
        if (Magic[0] != 0x9a)
            continue;
        Magic[1] = mt_usbtty_getc();
        if (Magic[1] != 0x23)
            continue;
        mt_usbtty_getcn(14, (char *)&Header->Type);
        Header->Magic = PACKET_HEADER_MAGIC;
        return;
    }
}

INT AgpReceiveData(UINTPTR Address, UINT32 BlockSize, UINT32 NumBlocks)
{
    PACKET_HEADER Header;
    CHAR *Buffer = (CHAR *)Address;
    UINT32 CurrentBlock = 0;
    while (1)
    {
        if (CurrentBlock == NumBlocks)
            break;
        AgpReceivePacketHeader(&Header);
        if (Header.Type == PT_DATA)
        {
            if (Header.DataLength != BlockSize)
            {
                UsbSendStatus(AG_STATUS_INVALID_TRANSFER_SIZE, NULL, 0);
                return 2;
            }
            mt_usbtty_getcn(BlockSize, Buffer);
            Buffer += BlockSize;
            CurrentBlock++;
        }
        else if (Header.Type == PT_ABORT)
            return 1;
    }

    return 0;
}

VOID AgpHandleCommand(UINT32 Command, UINT8 *Data, SIZE DataSize)
{
    static PACKET_HEADER Header;
    INT ReceiveStatus;

    UINT32 *Data32 = (UINT32 *)Data;
    UINT64 *Data64 = (UINT64 *)Data64;

    switch (Command)
    {
    case AG_CMD_REBOOT:
        UsbSendStatus(AG_STATUS_OK, NULL, 0);
        ResetBoard();
        break;
    case AG_CMD_ENABLE_CONSOLE:
        AgpConsoleEnabled = TRUE;
        UsbSendStatus(AG_STATUS_OK, NULL, 0);
        break;
    case AG_CMD_DISABLE_CONSOLE:
        AgpConsoleEnabled = FALSE;
        UsbSendStatus(AG_STATUS_OK, NULL, 0);
        break;
    case AG_CMD_PRINT_EFUSE_CFG:
    {
        UsbSendStatus(AG_STATUS_OK, NULL, 0);
        AgpPrintEfuseState();
        UsbSendStatus(AG_STATUS_OK, NULL, 0);
        break;
    }
    case AG_CMD_PRINT_BOOTINFO:
    {
        UsbSendStatus(AG_STATUS_OK, NULL, 0);
        AgpPrintBootInfo();
        UsbSendStatus(AG_STATUS_OK, NULL, 0);
        break;
    }
    case AG_CMD_SEND_LK:
    {
        if (DataSize < 8)
        {
            UsbSendStatus(AG_STATUS_INVALID_ARGUMENT, NULL, 0);
            return;
        }

        UINT32 BlockSize = ((UINT32 *)Data)[0];
        UINT32 NumBlocks = ((UINT32 *)Data)[1];
        if (BlockSize != 512 && BlockSize != 1024)
        {
            UsbSendStatus(AG_STATUS_INVALID_ARGUMENT, NULL, 0);
            return;
        }
        UsbSendStatus(AG_STATUS_OK, NULL, 0);

        ReceiveStatus = AgpReceiveData(LK_LOAD_ADDRESS, BlockSize, NumBlocks);
        if (ReceiveStatus == 0)
        {
            AgpHaveLk = TRUE;
            AgpLkSize = BlockSize * NumBlocks;
            UsbSendStatus(AG_STATUS_OK, NULL, 0);
        }
        else
        {
            AgpLkSize = 0;
            AgpHaveLk = FALSE;
        }

        break;
    }
    case AG_CMD_SEND_MTK_DA:
    {
        if (DataSize < 12)
        {
            UsbSendStatus(AG_STATUS_INVALID_ARGUMENT, NULL, 0);
            return;
        }

        UINT32 BlockSize = ((UINT32 *)Data)[0];
        UINT32 NumBlocks = ((UINT32 *)Data)[1];
        UINT32 DaSize = ((UINT32 *)Data)[2];
        if ((BlockSize != 512 && BlockSize != 1024) || DaSize > BlockSize * NumBlocks)
        {
            UsbSendStatus(AG_STATUS_INVALID_ARGUMENT, NULL, 0);
            return;
        }
        UsbSendStatus(AG_STATUS_OK, NULL, 0);

        ReceiveStatus = AgpReceiveData(MTK_DA_LOAD_ADDRESS, BlockSize, NumBlocks);
        if (ReceiveStatus == 0)
        {
            AgpHaveMtkDa = TRUE;
            AgpDaSize = BlockSize * NumBlocks;
            UsbSendStatus(AG_STATUS_OK, NULL, 0);
        }
        else
        {
            AgpDaSize = 0;
            AgpHaveMtkDa = FALSE;
        }

        break;
    }
    case AG_CMD_BOOT_LK:
        if (!AgpHaveLk)
        {
            UsbSendStatus(AG_STATUS_UNKNOWN_ERROR, NULL, 0);
        }
        else
        {
            AgpPatchLk((PVOID)LK_LOAD_ADDRESS, (PVOID)(LK_ENTRY_POINT), AgpLkSize, FASTBOOT, 0);
            UsbSendStatus(AG_STATUS_OK, NULL, 0);
            if (!AgpHaveAtf)
                AgpBootLk((PVOID)LK_LOAD_ADDRESS, AgpLkSize, LK_ENTRY_POINT, NULL, 0);
            else
                AgpBootLk((PVOID)LK_LOAD_ADDRESS, AgpLkSize, LK_ENTRY_POINT, (PVOID)ATF_LOAD_ADDRESS, AgpAtfSize);
        }
        break;
    case AG_CMD_BOOT_MTK_DA:
    {
        UINT32 Flags = 0;
        if (DataSize >= 4)
        {
            Flags = ((UINT32 *)Data)[0];
        }
        if ((Flags & AG_BOOT_MTK_DA_USE_LK && !AgpHaveLk) || !AgpHaveMtkDa)
        {
            UsbSendStatus(AG_STATUS_UNKNOWN_ERROR, NULL, 0);
            return;
        }
        UsbSendStatus(AG_STATUS_OK, NULL, 0);
        AgpBootMtkDa(MTK_DA_LOAD_ADDRESS, AgpDaSize, LK_ENTRY_POINT, !!(Flags & AG_BOOT_MTK_DA_USE_LK));

        break;
    }
    case AG_CMD_SEND_ATF:
    {
        if (DataSize < 8)
        {
            UsbSendStatus(AG_STATUS_INVALID_ARGUMENT, NULL, 0);
            return;
        }

        UINT32 BlockSize = ((UINT32 *)Data)[0];
        UINT32 NumBlocks = ((UINT32 *)Data)[1];

        if (BlockSize != 512 && BlockSize != 1024)
        {
            UsbSendStatus(AG_STATUS_INVALID_ARGUMENT, NULL, 0);
            return;
        }
        if (BlockSize * NumBlocks > ATF_MAX_SIZE)
        {
            UsbSendStatus(AG_STATUS_TOO_BIG, NULL, 0);
            return;
        }
        UsbSendStatus(AG_STATUS_OK, NULL, 0);

        ReceiveStatus = AgpReceiveData(ATF_LOAD_ADDRESS, BlockSize, NumBlocks);
        if (ReceiveStatus == 0)
        {
            AgpHaveAtf = TRUE;
            AgpAtfSize = BlockSize * NumBlocks;
            UsbSendStatus(AG_STATUS_OK, NULL, 0);
        }
        else
        {
            AgpAtfSize = 0;
            AgpHaveAtf = FALSE;
        }

        break;
    }
#if 0
    case AG_CMD_SEND_TEE:
    {
        if (DataSize < 8)
        {
            UsbSendStatus(AG_STATUS_INVALID_ARGUMENT, NULL, 0);
            return;
        }

        UINT32 BlockSize = ((UINT32 *)Data)[0];
        UINT32 NumBlocks = ((UINT32 *)Data)[1];
        if (BlockSize != 512 && BlockSize != 1024)
        {
            UsbSendStatus(AG_STATUS_INVALID_ARGUMENT, NULL, 0);
            return;
        }
        UsbSendStatus(AG_STATUS_OK, NULL, 0);

        CHAR *Buffer = (CHAR *)TEE_LOAD_ADDRESS;
        UINT32 CurrentBlock = 0;
        while (1)
        {
            if (CurrentBlock == NumBlocks)
                break;
            AgpReceivePacketHeader(&Header);
            if (Header.Type == PT_DATA)
            {
                if (Header.DataLength != BlockSize)
                {
                    UsbSendStatus(AG_STATUS_INVALID_TRANSFER_SIZE, NULL, 0);
                    return;
                }
                mt_usbtty_getcn(BlockSize, Buffer);
                Buffer += BlockSize;
                CurrentBlock++;
            }
            else if (Header.Type == PT_ABORT)
                return;
        }
        UsbSendStatus(AG_STATUS_OK, NULL, 0);
        break;
    }
#endif
    case AG_CMD_GET_LOGS:
    {
        UINT32 Flags = 0;
        if (DataSize > 4)
            Flags = ((UINT32 *)Data)[0];
        UsbSendStatus(AG_STATUS_OK, NULL, 0);
        AgpRetrieveRamConsole(Flags);
        UsbSendStatus(AG_STATUS_OK, NULL, 0);
        break;
    }
    case AG_CMD_MMC_INIT:
    {
        if (DataSize < 4)
        {
            RtlPrintf("AG_CMD_MMC_INIT: missing required arguments\n");
            UsbSendStatus(AG_STATUS_INVALID_ARGUMENT, NULL, 0);
            return;
        }
        UINT32 Index = ((UINT32 *)Data)[0];
        if (Index >= MtkGetMmcControllerCount())
        {
            RtlPrintf("AG_CMD_MMC_INIT: index exceeds controller count, %d >= %d\n", Index, MtkGetMmcControllerCount());
            UsbSendStatus(AG_STATUS_INVALID_ARGUMENT, NULL, 0);
            return;
        }
        if (N1_SUCCESS(MtkMmcInit(Index)))
        {
            AgpReadPartitionTable(Index);
            UsbSendStatus(AG_STATUS_OK, NULL, 0);
        }
        else
            UsbSendStatus(AG_STATUS_UNKNOWN_ERROR, NULL, 0);
        break;
    }
    case AG_CMD_MMC_INFO:
    {
        UINT32 Info[1] = {MtkGetMmcControllerCount()};
        UsbSendStatus(AG_STATUS_OK, Info, sizeof(Info));
        break;
    }
    case AG_CMD_READ_P:
    case AG_CMD_WRITE_P:
    {
        if (DataSize <= 20)
        {
            UsbSendStatus(AG_STATUS_INVALID_ARGUMENT, NULL, 0);
            return;
        }
        UINT64 Offset = *(UINT64 *)(Data + 0);
        UINT64 Length = *(UINT64 *)(Data + 8);
        UINT32 Mmc = *(UINT64 *)(Data + 16);
        CONST CHAR *PartitionName = (CONST CHAR *)(Data + 20);
        if (Data[DataSize] != '\0')
        {
            RtlPrintf("string is not null terminated\n");
            UsbSendStatus(AG_STATUS_INVALID_ARGUMENT, NULL, 0);
            return;
        }

        PAG_PTE Pt = AgpLookupPartition(Mmc, PartitionName);
        if (!Pt)
        {
            UsbSendStatus(AG_STATUS_NO_SUCH_PARTITION, NULL, 0);
            return;
        }

        if ((!(Pt->Policy & AG_P_SECPOL_ALLOW_READ) && Command == AG_CMD_READ_P) || (!(Pt->Policy & AG_P_SECPOL_ALLOW_WRITE) && Command == AG_CMD_WRITE_P))
        {
            UsbSendStatus(AG_STATUS_ACCESS_DENIED, NULL, 0);
            return;
        }

        UINT64 MaxLength = (Pt->LastLba - Pt->FirstLba + 1) * MMC_BLOCK_SIZE;
        if (Length == 0xFFFFFFFF && Command == AG_CMD_READ_P)
        {
            Length = MaxLength;
        }
        else
        {
            if (Length > MaxLength)
            {
                UsbSendStatus(AG_STATUS_INVALID_ARGUMENT, NULL, 0);
                return;
            }
        }

        if (Length % MMC_BLOCK_SIZE != 0 || Offset != 0)
        {
            RtlPrintf("unimplemented\n");
            UsbSendStatus(AG_STATUS_INVALID_ARGUMENT, NULL, 0);
            return;
        }

        char Buffer[MMC_BLOCK_SIZE];
        if (Command == AG_CMD_READ_P)
        {
            RtlPrintf("reading data from partition %s offset=0x%016llx size=0x%016llx mmc=%d\n", PartitionName, Offset, Length, Mmc);
            UsbSendStatus(AG_STATUS_OK, &Length, sizeof(Length));

            UINT64 Left = Length;

            INT Block = 0;
            while (Left > 0)
            {
                PACKET_HEADER Header;
                Header.Magic = PACKET_HEADER_MAGIC;
                Header.Type = PT_DATA;
                Header.DataLength = MIN(Left, MMC_BLOCK_SIZE);
                Header.Serial = 0;

                mt_usbtty_putcn(PACKET_HEADER_SIZE, (char *)&Header, 1);
                MtkMmcRead(Mmc, Buffer, 1, Pt->FirstLba + Block);
                //mt_usbtty_putcn(Header.DataLength, Buffer, 1);
                mt_usbtty_puts((const char *)Buffer, Header.DataLength);
                Left -= Header.DataLength;
                Block++;
            }
            UsbSendStatus(AG_STATUS_OK, NULL, 0);
        }
        else if (Command == AG_CMD_WRITE_P)
        {
            RtlPrintf("writing data to partition %s offset=0x%016llx size=0x%016llx mmc=%d\n", PartitionName, Offset, Length, Mmc);
            UsbSendStatus(AG_STATUS_OK, NULL, 0);

            UINT64 Left = Length;
            INT Block = 0;
            while (Left > 0)
            {
                INT r = AgpReceiveData((UINTPTR)Buffer, MMC_BLOCK_SIZE, 1);
                if (r != 0)
                    return;
                MtkMmcWrite(Mmc, Buffer, 1, Pt->FirstLba + Block);
                Block++;
                Left -= MMC_BLOCK_SIZE;
            }
            UsbSendStatus(AG_STATUS_OK, NULL, 0);
        }

        break;
    }
    case AG_CMD_SECPOL_DUMP:
    {
        AgpDumpSecpol();
        UsbSendStatus(AG_STATUS_OK, NULL, 0);
        break;
    }
    case AG_CMD_SECPOL_LOCK:
        AgpLockDevice();
        UsbSendStatus(AG_STATUS_OK, NULL, 0);
        break;
    case AG_CMD_SECPOL_UNLOCK:
        AgpUnlockDevice();
        UsbSendStatus(AG_STATUS_OK, NULL, 0);
        break;
    case AG_CMD_RUN_TEST:
    {
        if (DataSize < 4)
        {
            UsbSendStatus(AG_STATUS_INVALID_ARGUMENT, NULL, 0);
            return;
        }
        AgpRunTest(Data32[0]);
        UsbSendStatus(AG_STATUS_OK, NULL, 0);
        break;
    }
    case AG_CMD_SEND_PL:
    {
        if (DataSize < 8)
        {
            UsbSendStatus(AG_STATUS_INVALID_ARGUMENT, NULL, 0);
            return;
        }

        UINT32 BlockSize = ((UINT32 *)Data)[0];
        UINT32 NumBlocks = ((UINT32 *)Data)[1];
        if (BlockSize != 512 && BlockSize != 1024)
        {
            UsbSendStatus(AG_STATUS_INVALID_ARGUMENT, NULL, 0);
            return;
        }
        UsbSendStatus(AG_STATUS_OK, NULL, 0);

        ReceiveStatus = AgpReceiveData(PL_LOAD_ADDRESS, BlockSize, NumBlocks);
        if (ReceiveStatus == 0)
        {
            AgpHavePl = TRUE;
            UsbSendStatus(AG_STATUS_OK, NULL, 0);
        }
        else
        {
            AgpHavePl = FALSE;
        }
        break;
    }
    case AG_CMD_BOOT_PL:
    {
        if (!AgpHavePl)
        {
            UsbSendStatus(AG_STATUS_UNKNOWN_ERROR, NULL, 0);
        }
        else
        {
            UsbSendStatus(AG_STATUS_OK, NULL, 0);
            AgpBootPl(PL_LOAD_ADDRESS);
        }
        break;
    }
    case AG_CMD_ENTER_BOOTROM:
        RtlPrintf("Entering Boot ROM mode ...");
        UsbSendStatus(AG_STATUS_OK, NULL, 0);
        AgEnterBootROM();
        break;
    default:
        UsbSendStatus(AG_STATUS_UNSUPPORTED_COMMAND, NULL, 0);
        break;
    }
}

VOID AgpProcessPacket(PPACKET_HEADER Header, UINT8 *Buffer)
{
    if (Header->Type == PT_COMMAND)
    {
        if (Header->DataLength < 4)
            UsbSendStatus(AG_STATUS_INVALID_ARGUMENT, NULL, 0);
        else
        {
            AgpHandleCommand(((UINT32 *)Buffer)[0], Buffer + 4, Header->DataLength - 4);
        }
    }
    else
    {
        if (Header->Type == PT_ABORT)
        {
        }
        else
            UsbSendStatus(AG_STATUS_UNKNOWN_ERROR, NULL, 0);
    }
}

extern VOID _jump(UINT32 Address, UINT32 Arg1, UINT32 Arg2) __attribute__((noreturn));

VOID N1Main(UINT32 BootArg, UINT32 BootArgLength)
{
    RtlFillMemory(__bss_start, 0, (UINTPTR)__bss_end - (UINTPTR)__bss_start);

    mtk_timer_init();
    usbdl_init();
    udelay(1000);
    usb_disconnect();
    if (usb_connect(1000) == FALSE)
    {
        ResetBoard();
    }

    AgpGlobalHeap = RtlCreateHeap(AgpGlobalHeapPool, GLOBAL_HEAP_POOL_SIZE);

    BOOLEAN SetupComplete = FALSE;
    PACKET_HEADER Header;

    while (1)
    {
        AgpReceivePacketHeader(&Header);

        if (Header.Type != PT_SETUP && !SetupComplete)
        {
            UsbSendStatus(AG_STATUS_SETUP_REQUIRED, NULL, 0);
            continue;
        }

        if (Header.Type == PT_SETUP)
        {
            AgpSetup();
            SetupComplete = TRUE;
            continue;
        }

        if (Header.DataLength > BUFFER_SIZE)
        {
            UsbSendStatus(AG_STATUS_TOO_MUCH_DATA, NULL, 0);
            continue;
        }

        RtlZeroMemory(Buffer, BUFFER_SIZE);
        mt_usbtty_getcn(Header.DataLength, Buffer);
        AgpProcessPacket(&Header, (UINT8 *)Buffer);
    }

    ResetBoard();

    while (1)
    {
    }

#if 0
    VOID DisableSec(UINT32 Address)
    {
        WRITE32(Address, 0xe3a00000);     // mov r0, #0
        WRITE32(Address + 4, 0xe12fff1e); // bx lr
    }

    // disable SBC
    DisableSec(0x00220DB0);
    // disable DAA
    DisableSec(0x00220E40);
    // disable SLA
    DisableSec(0x00220E8C);

    WRITE32(0x00102080, 0x3B6C243C);
    WRITE32(0x00102084, 0xF843E0A);

    __asm__ __volatile__("" ::
                             : "memory");
#endif
}
